package com.xfsg.core.rocketmq;

import com.xfsg.core.rocketmq.hook.ConsumerLogTraceHook;
import com.xfsg.core.rocketmq.hook.ProducerLogTraceHook;
import lombok.extern.slf4j.Slf4j;
import org.apache.rocketmq.spring.core.RocketMQTemplate;
import org.apache.rocketmq.spring.support.DefaultRocketMQListenerContainer;
import org.springframework.beans.BeansException;
import org.springframework.boot.CommandLineRunner;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;

/**
 * @author ZhangYi
 * @classname RocketMqRunner
 * @description TODO
 * @date 2022/2/18 22:58
 */
@Slf4j
public class RocketMqRunner  implements ApplicationContextAware, CommandLineRunner {

    private ApplicationContext applicationContext;

    private String serviceId;

    public RocketMqRunner() {
    }

    public RocketMqRunner(String serviceId) {
        this.serviceId = serviceId;
    }

    @Override
    public void run(String... args) throws Exception {
        for (String beanName : applicationContext.getBeanNamesForType(RocketMQTemplate.class)) {
            RocketMQTemplate bean = applicationContext.getBean(beanName, RocketMQTemplate.class);
            if (null != bean.getProducer()) {
                log.info("producer={} register message hook", beanName);
                bean.getProducer().getDefaultMQProducerImpl().registerSendMessageHook(new ProducerLogTraceHook(serviceId));
            }
        }
        for (String beanName : applicationContext.getBeanNamesForType(DefaultRocketMQListenerContainer.class)) {
            DefaultRocketMQListenerContainer bean = applicationContext.getBean(beanName, DefaultRocketMQListenerContainer.class);
            if (null != bean.getConsumer()) {
                log.info("consumer={} register message hook", beanName);
                bean.getConsumer().getDefaultMQPushConsumerImpl().registerConsumeMessageHook(new ConsumerLogTraceHook());
            }
        }
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

}
